#!/Pub/Users/wangyk/software/molecular_dockinig/vina/.pixi/envs/default/bin/python3
import argparse
import json
import sys
from pathlib import Path
import subprocess

def validate_json_data(data):
    """验证JSON数据结构"""
    if not all(key in data for key in ["center", "size"]):
        raise ValueError("JSON缺少'center'或'size'字段")
    if len(data["center"]) != 3 or len(data["size"]) != 3:
        raise ValueError("'center'或'size'必须是长度为3的数组")
    return True

def main(args):
    try:
        # 使用Path处理路径更安全
        ligand_path = Path(args.ligand_pdb)
        json_path = Path(args.center_and_autobox_info)
        
        # ligand_path = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/R/Convertional_Transcriptome/Molecular_Docking_cmd/v1.1/test/CTSS_1glo_no_water.pdb"
        
        # 尝试纠正HIS问题
        # import os
        # os.system(f"cp {ligand_path} {ligand_path}.bak")
        
        # with open(f"{ligand_path}.bak") as f_in, open(ligand_path, "w") as f_out:
        #     for line in f_in:
        #         if line.startswith("ATOM") and "HIS" in line:
        #             line = line.replace("HIS", "HID")  # 或 HIE/HIP
        #         f_out.write(line)
        
        # 验证文件存在
        if not ligand_path.exists():
            raise FileNotFoundError(f"配体文件不存在: {ligand_path}")
        if not json_path.exists():
            raise FileNotFoundError(f"JSON配置文件不存在: {json_path}")

        # 读取JSON数据
        with open(json_path, 'r') as f:
            data = json.load(f)
            validate_json_data(data)
            
            center = data["center"]
            size = data["size"]
            print(f"加载配置: center={center}, size={size}")

        # 准备输出路径
        output_path = ligand_path.with_name(f"{ligand_path.stem}_prepared")
        
        # 使用subprocess.run更安全的执行命令
        mk_prepare_receptor_src = "/Pub/Users/wangyk/software/molecular_dockinig/vina/.pixi/envs/default/bin/mk_prepare_receptor.py"
        
        cmd = [
            mk_prepare_receptor_src,
            '-i', str(ligand_path),
            '-o', str(output_path),
            '-p', 
            '-v',
            '--box_size', str(size[0]), str(size[1]), str(size[2]),
            '--box_center', str(center[0]), str(center[1]), str(center[2]),
            '--default_altloc','A'
        ]
        
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        print(result.stdout)
        
    except subprocess.CalledProcessError as e:
        print(f"命令执行失败: {e.stderr}", file=sys.stderr)
        sys.exit(e.returncode)
    except Exception as e:
        print(f"错误: {str(e)}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="准备对接分子受体并设置对接盒子参数",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("ligand_pdb", help="输入PDB文件路径")
    parser.add_argument("center_and_autobox_info", help="包含中心坐标和盒子尺寸的JSON文件")
    
    args = parser.parse_args()
    main(args)